#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import sys
import oss2
import time
# sys.path.append('/opt/work/caffe/python')
sys.path.insert(0, '.')

import os
import argparse
import multiprocessing

############################################################################


def download_run(file_queue, out_queue, id, args):
    # 阿里云主账号AccessKey拥有所有API的访问权限，风险很高。强烈建议您创建并使用RAM账号进行API访问或日常运维，请登录 https://ram.console.aliyun.com 创建RAM账号。
    auth = oss2.Auth('LTAIxABOHnLXUh8P', 'gUaZTagYxFzh6k6Iu7hk0HvuTEOSlZ')
    # Endpoint以杭州为例，其它Region请按实际情况填写。
    bucket = oss2.Bucket(auth, 'http://oss-cn-shenzhen.aliyuncs.com', 'erised-screenshots')
    output_dir = args.output_dir
    try:
        while True:
            file_info = file_queue.get(timeout=5)
            if file_info is None:
                break
            key_file_name = file_info
            video_download_path = os.path.join(output_dir, key_file_name)
            dirname = os.path.dirname(video_download_path)
            try:
                if not os.path.exists(dirname):
                    os.makedirs(dirname)
            except Exception as e:
                pass
            try:
                bucket.get_object_to_file(key_file_name, video_download_path)
            except Exception as e:
                print('Download Error:', video_download_path)
                print(e)
                continue

            while out_queue.qsize() > 100:
                time.sleep(0.01)
            out_queue.put((key_file_name, False, id))
    except Exception as e:
        if str(e) != '':
            print(e)

    out_queue.put((None, True, id))


def out_run(out_queue, args, total_file_count, output_dir):
    file_count = 0
    out_file_dir = os.path.dirname(args.download_video_file)
    if not os.path.exists(out_file_dir):
        os.makedirs(out_file_dir)
    out_file = open(args.download_video_file, 'a')
    try:
        finish_worker_count = 0
        while True:
            file_info = out_queue.get(block=True)
            if file_info is None:
                break
            file_name, finish, id = file_info
            if finish:
                print('Proc{} finish'.format(id, ))
                finish_worker_count += 1
                if args.proccess_count <= finish_worker_count:
                    break
                continue
            file_count += 1
            out_info = '{}\n'.format(file_name)
            out_file.write(out_info)
            print('{:06f}, Proc{}, Count: {}/{}, {}'.format(time.time(), id, file_count, total_file_count, file_name))
            if file_count % 20 == 0:
                out_file.flush()
    except Exception as e:
        print(e)
    out_file.close()
    os.system('chmod a+wr {}'.format(args.download_video_file))


def main():
    parser = argparse.ArgumentParser(description="Video Download")
    parser.add_argument(
        "--output_dir", default="/rootfs/media/kasim/DataSet/Record2", help="path to download video dir", type=str
    )
    parser.add_argument(
        "--video_list", default="./server_video_list.txt", help="path to video file list", type=str
    )
    parser.add_argument(
        "--download_video_file", default="./download_video_list.txt", help="path to download video file file", type=str
    )
    parser.add_argument('--proccess_count', type=int, default=4, help='download proccess count')
    args = parser.parse_args()
    output_dir = args.output_dir
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
        os.system('chmod a+wr {}'.format(output_dir))

    download_videolist = set()
    if os.path.exists(args.download_video_file):
        with open(args.download_video_file, 'r') as file:
            for file_name in file.readlines():
                download_videolist.add(file_name.split(',')[0].strip())

    need_download_videolist = []
    file_count = 0
    with open(args.video_list, 'r') as file:
        for file_info in file.readlines():
            file_name = file_info.split(',')[0].strip()
            if file_name in download_videolist:
                continue
            need_download_videolist.append(file_name)
            file_count += 1
            if file_count % 100 == 0:
                print('Need Download File Count:', file_count)
        print('Need Download File Count:', file_count)

    file_queue = multiprocessing.Queue()
    out_queue = multiprocessing.Queue()

    workers = []
    for i in range(args.proccess_count):
        workers.append(multiprocessing.Process(target=download_run, args=(file_queue, out_queue, i, args)))

    for i in range(args.proccess_count):
        workers[i].start()

    wait_count_thr = args.proccess_count * 10
    out_worker = multiprocessing.Process(target=out_run, args=(out_queue, args, file_count, output_dir))
    out_worker.start()

    for need_download_video in need_download_videolist:
        while file_queue.qsize() > wait_count_thr:
            time.sleep(0.01)
        file_queue.put(need_download_video)

    out_worker.join()

    for i in range(args.proccess_count):
        workers[i].join()


if __name__ == '__main__':
    main()
